﻿' (C) Copyright 2009-2010 by Autodesk, Inc. 
'
' Permission to use, copy, modify, and distribute this software in
' object code form for any purpose and without fee is hereby granted, 
' provided that the above copyright notice appears in all copies and 
' that both that copyright notice and the limited warranty and
' restricted rights notice below appear in all supporting 
' documentation.
'
' AUTODESK PROVIDES THIS PROGRAM "AS IS" AND WITH ALL FAULTS. 
' AUTODESK SPECIFICALLY DISCLAIMS ANY IMPLIED WARRANTY OF
' MERCHANTABILITY OR FITNESS FOR A PARTICULAR USE.  AUTODESK, INC. 
' DOES NOT WARRANT THAT THE OPERATION OF THE PROGRAM WILL BE
' UNINTERRUPTED OR ERROR FREE.
'
' Use, duplication, or disclosure by the U.S. Government is subject to 
' restrictions set forth in FAR 52.227-19 (Commercial Computer
' Software - Restricted Rights) and DFAR 252.227-7013(c)(1)(ii)
' (Rights in Technical Data and Computer Software), as applicable.
'

Imports System.Collections.Generic
Imports System.Reflection
Imports System.Resources
Imports System
Imports Microsoft.Win32
Imports Autodesk.AutoCAD.DatabaseServices
Imports Autodesk.AutoCAD.Runtime

Namespace DemandLoading
  Public Class RegistryUpdate
        Public Shared Sub RegisterForDemandLoading(Optional ByVal p_force As Boolean = False)
            ' Get the assembly, its name and location

            Dim assem As Assembly = Assembly.GetExecutingAssembly()
            Dim name As String = assem.GetName().Name
            Dim path As String = assem.Location

            ' We'll collect information on the commands
            ' (we could have used a map or a more complex
            ' container for the global and localized names
            ' - the assumption is we will have an equal
            ' number of each with possibly fewer groups)

            Dim globCmds As New List(Of String)()
            Dim locCmds As New List(Of String)()
            Dim groups As New List(Of String)()

            ' Iterate through the modules in the assembly

            Dim mods As [Module]() = assem.GetModules(True)
            For Each [mod] As [Module] In mods
                ' Within each module, iterate through the types

                Dim types As Type() = [mod].GetTypes()
                For Each type As Type In types
                    ' We may need to get a type's resources

                    Dim rm As New ResourceManager(type.FullName, assem)
                    rm.IgnoreCase = True

                    ' Get each method on a type

                    Dim meths As MethodInfo() = type.GetMethods()
                    For Each meth As MethodInfo In meths
                        ' Get the methods custom command attribute(s)

                        Dim attbs As Object() = _
                          meth.GetCustomAttributes( _
                            GetType(CommandMethodAttribute), True)
                        For Each attb As Object In attbs
                            Dim cma As CommandMethodAttribute = _
                              TryCast(attb, CommandMethodAttribute)
                            If cma IsNot Nothing Then
                                ' And we can finally harvest the information
                                ' about each command

                                Dim globName As String = cma.GlobalName
                                Dim locName As String = globName
                                Dim lid As String = cma.LocalizedNameId

                                ' If we have a localized command ID,
                                ' let's look it up in our resources

                                If lid IsNot Nothing Then
                                    ' Let's put a try-catch block around this
                                    ' Failure just means we use the global
                                    ' name twice (the default)

                                    Try
                                        locName = rm.GetString(lid)
                                    Catch
                                    End Try
                                End If

                                ' Add the information to our data structures

                                globCmds.Add(globName)
                                locCmds.Add(locName)

                                If cma.GroupName IsNot Nothing AndAlso _
                                  Not groups.Contains(cma.GroupName) Then
                                    groups.Add(cma.GroupName)
                                End If
                            End If
                        Next
                    Next
                Next
            Next

            ' Let's register the application to load on demand (12)
            ' if it contains commands, otherwise we will have it
            ' load on AutoCAD startup (2)

            'Dim flags As Integer = (If(globCmds.Count > 0, 12, 2))
            Dim flags As Integer
            If p_force Then
                flags = 2
            Else
                flags = (If(globCmds.Count > 0, 12, 2))
            End If

            ' By default let's create the commands in HKCU
            ' (pass false if we want to create in HKLM)

            CreateDemandLoadingEntries(name, path, globCmds, locCmds, _
              groups, flags, True)
        End Sub

    Public Shared Sub UnregisterForDemandLoading()
      RemoveDemandLoadingEntries(True)
    End Sub

    ' Helper functions

    Private Shared Sub CreateDemandLoadingEntries( _
      ByVal name As String, ByVal path As String, _
      ByVal globCmds As List(Of String), _
      ByVal locCmds As List(Of String), _
      ByVal groups As List(Of String), _
      ByVal flags As Integer, _
      ByVal currentUser As Boolean)

      ' Choose a Registry hive based on the function input

      Dim hive As RegistryKey = _
        If(currentUser, Registry.CurrentUser, Registry.LocalMachine)

      ' Open the main AutoCAD (or vertical) and "Applications" keys

      Dim ack As RegistryKey = _
        hive.OpenSubKey( _
          HostApplicationServices.Current.RegistryProductRootKey, _
          True)
      Using ack
        Dim appk As RegistryKey = ack.CreateSubKey("Applications")
        Using appk

          ' Already registered? Just return

          Dim subKeys As String() = appk.GetSubKeyNames()
          For Each subKey As String In subKeys
            If subKey.Equals(name) Then
              appk.Close()
              Exit Sub
            End If
          Next

          ' Create the our application's root key and its values

          Dim rk As RegistryKey = appk.CreateSubKey(name)
          Using rk
            rk.SetValue( _
              "DESCRIPTION", name, RegistryValueKind.[String])
            rk.SetValue("LOADCTRLS", flags, RegistryValueKind.DWord)
            rk.SetValue("LOADER", path, RegistryValueKind.[String])
            rk.SetValue("MANAGED", 1, RegistryValueKind.DWord)

            ' Create a subkey if there are any commands...

            If (globCmds.Count = locCmds.Count) _
              AndAlso globCmds.Count > 0 Then
              Dim ck As RegistryKey = rk.CreateSubKey("Commands")
              Using ck
                For i As Integer = 0 To globCmds.Count - 1
                  ck.SetValue(globCmds(i), locCmds(i), _
                    RegistryValueKind.[String])
                Next
              End Using
            End If

            ' And the command groups, if there are any

            If groups.Count > 0 Then
              Dim gk As RegistryKey = rk.CreateSubKey("Groups")
              Using gk
                For Each grpName As String In groups
                  gk.SetValue(grpName, grpName, _
                    RegistryValueKind.[String])
                Next
              End Using
            End If
          End Using
        End Using
      End Using
    End Sub

    Private Shared Sub RemoveDemandLoadingEntries( _
      ByVal currentUser As Boolean)

      Try

        ' Choose a Registry hive based on the function input

        Dim hive As RegistryKey = _
          If(currentUser, Registry.CurrentUser, Registry.LocalMachine)

        ' Open the main AutoCAD (or vertical) and "Applications" keys

        Dim ack As RegistryKey = _
          hive.OpenSubKey( _
            HostApplicationServices.Current.RegistryProductRootKey)
        Dim appk As RegistryKey = _
          ack.OpenSubKey("Applications", True)

        ' Delete the key with the same name as this assembly

        appk.DeleteSubKeyTree( _
          Assembly.GetExecutingAssembly().GetName().Name)
        appk.Close()

      Catch
      End Try
    End Sub
  End Class
End Namespace
